from flask import Flask, request, redirect, url_for, render_template
import os
import constant
import cv2 as cv
from utils import get_img_name, result_util, mysql_PooledDB
from PIL import Image
import numpy as np
import threading

app = Flask(__name__, static_folder='', static_url_path='')
R = result_util()
# 初始化数据连接池
mysql = mysql_PooledDB()
# 线程锁
lock = threading.Lock()


@app.route('/')
def index():
    return render_template('index.html')


def getImageAndLabels(path):
    faceSamples = []
    ids = []
    image_paths = [os.path.join(path, f) for f in os.listdir(path)]
    # 检测人脸
    face_detect = cv.CascadeClassifier(constant.CASCADE_CLASSIFIER_FILE)

    for img_path in image_paths:
        # 打开图片
        PIL_img = Image.open(img_path).convert('L')
        img_numpy = np.array(PIL_img, 'uint8')
        faces = face_detect.detectMultiScale(img_numpy)
        file_name = get_img_name(is_new=False, img_path=img_path)
        sql = "SELECT `id` FROM `file_info` WHERE file_name = %s"
        res = mysql.select(sql=sql, param=file_name)
        for x, y, w, h in faces:
            faceSamples.append(img_numpy[y:y + h, x:x + w])
            ids.append(res[0]['id'])

    return faceSamples, ids


# 机器学习模块
@app.route('/admin/learn', methods=['POST'])
def learn():
    # 图片文件夹路径
    data_path = constant.ADMIN_LEARN_UPLOAD_PATH
    # 获取图像数组和id标签
    faces, ids = getImageAndLabels(data_path)
    recognizer = cv.face.LBPHFaceRecognizer_create()
    recognizer.train(faces, np.array(ids))
    # 保存文件
    recognizer.write(os.path.join(constant.TRAINER_PATH, 'trainer.yml'))
    return R.ok()


@app.route('/admin/uploadList', methods=['POST'])
def upload_file_list():
    # 获取锁 由于多文件上传是多线程的，不加锁数据库插入会错乱
    lock.acquire()
    file = request.files['file']
    person_name = request.form['personName']
    filename = get_img_name(file_name=file.filename)
    filepath = os.path.join(constant.ADMIN_LEARN_UPLOAD_PATH, filename)
    upload_path = os.path.join(constant.RELATIVE_ADMIN_LEARN_UPLOAD_PATH, filename)
    file.save(filepath)
    sql = "INSERT INTO `file_info` (`file_name`, `file_url`, `person_name`) VALUES (%s, %s, %s)"
    data = (filename, upload_path, person_name)
    try:
        mysql.insert(sql=sql, data=data)
    finally:
        # 释放锁
        lock.release()
    return R.ok()


@app.route('/admin')
def admin():
    return render_template('admin/index.html')


@app.route('/admin/upload')
def upload():
    return render_template('admin/upload.html')


@app.route('/admin/learn', methods=['GET'])
def learn_page():
    return render_template('admin/learn.html')


def face_detect_fun(img):
    # 图片设置成灰度
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # 加载特征数据
    face_detect = cv.CascadeClassifier(constant.CASCADE_CLASSIFIER_FILE)
    face = face_detect.detectMultiScale(gray)
    for x, y, w, h in face:
        cv.rectangle(img, (x, y), (x + w, y + h), color=(0, 0, 255), thickness=2)


# 人脸检测
@app.route('/detect', methods=['POST'])
def detect():
    img_path = request.form['imgPath']  # 得到图片路径字符串
    # 加载图片
    result_img = cv.imread(os.path.join(constant.CURRENT_PATH, img_path))
    # 调用人脸检测方法
    face_detect_fun(result_img)
    filename = get_img_name(is_new=False, img_path=img_path)
    # 写入文件
    cv.imwrite(os.path.join(constant.DETECT_RESULT_PATH, filename), result_img)
    return R.ok(data=os.path.join(constant.RELATIVE_DETECT_RESULT_PATH, filename))


# 人脸识别
@app.route('/discern', methods=['POST'])
def discern():
    # 加载训练数据集文件
    recognizer = cv.face.LBPHFaceRecognizer_create()
    recognizer.read(os.path.join(constant.TRAINER_PATH, 'trainer.yml'))
    img_path = request.form['imgPath']  # 得到图片路径字符串
    # 准备识别图片
    img = cv.imread(os.path.join(constant.CURRENT_PATH, img_path))
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # 加载特征数据
    face_discern = cv.CascadeClassifier(constant.CASCADE_CLASSIFIER_FILE)
    face = face_discern.detectMultiScale(gray)
    person_name = '无法识别'
    msg = ''
    for x, y, w, h in face:
        # 人脸识别
        file_id, confidence = recognizer.predict(gray[y:y + h, x:x + w])
        print('id', file_id, '置信评分', confidence)
        if confidence < 80:
            sql = "SELECT `person_name` FROM `file_info` WHERE id = %s"
            res = mysql.select(sql=sql, param=file_id)
            person_name = res[0]['person_name']
            cv.rectangle(img, (x, y), (x + w, y + h), color=(0, 255, 0), thickness=2)
            msg = " 置信评分：" + str(100 - int(confidence)) + "%"
    return R.ok(data="# &nbsp;" + person_name + "&nbsp; #" + msg)


@app.route('/detect/upload', methods=['POST'])
def detect_upload_file():
    file = request.files['file']
    filename = get_img_name(file_name=file.filename)
    filepath = os.path.join(constant.DETECT_UPLOAD_PATH, filename)
    file.save(filepath)
    upload_path = os.path.join(constant.RELATIVE_DETECT_UPLOAD_PATH, filename)
    return R.ok(data=upload_path)


@app.route('/discern/upload', methods=['POST'])
def discern_upload_file():
    file = request.files['file']
    filename = get_img_name(file_name=file.filename)
    filepath = os.path.join(constant.DISCERN_UPLOAD_PATH, filename)
    file.save(filepath)
    upload_path = os.path.join(constant.RELATIVE_DISCERN_UPLOAD_PATH, filename)
    return R.ok(data=upload_path)


if __name__ == '__main__':
    app.run()
